# -*- coding: utf-8 -*-

import logging
import subprocess
import re
import time
from datetime import datetime
from dingtalkchatbot.chatbot import DingtalkChatbot

# 预定义的IP列表
allowed_ips = {"10.210.2.45", "10.200.1.115"}

# 设置 Webhook 地址
webhook = "https://oapi.dingtalk.com/robot/send?access_token=c54fa6837d584b0abc56c0a7d0456ca37d59213dcc7ea3f6f104b20de3c71284"
secret = "SEC9de9f15e7d9bc3be3c0c5f50adcf283ff26189940919486b9cdb0feda9ee45b1"
# 初始化机器人
bot = DingtalkChatbot(webhook=webhook, secret=secret)

# 配置日志记录
logging.basicConfig(
    filename="/data/py/transfer.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
)


def nowTime():
    now = datetime.now()
    # 将时间格式化为字符串
    formatted_time = now.strftime("%Y-%m-%d %H:%M:%S")
    return formatted_time


def get_connected_ips():
    # 执行命令获取当前SSH连接
    cmd = "netstat -atnp | grep 'ESTABLISHED.*sshd' | awk '{print $5}'"
    result = subprocess.run(cmd, shell=True, stdout=subprocess.PIPE, text=True)
    output = result.stdout

    # 使用正则表达式解析IP地址
    ip_regex = re.compile(r"\d+\.\d+\.\d+\.\d+")
    connected_ips = set(ip_regex.findall(output))

    return connected_ips


def send_notification(messages, ip):
    current_time = nowTime()
    try:
        # 发送请求
        bot.send_text(messages)
        warn_message = "{}，异常 IP 地址 {} 通过 ssh 访问服务器。".format(current_time, ip)
        print(messages)
        logging.warning(warn_message)
    except Exception as e:
        error_message = "{}，发送钉钉通知时发生异常：{}。".format(current_time, e)
        print(error_message)
        logging.error(error_message)


def main():
    while True:
        connected_ips = get_connected_ips()
        current_time = nowTime()
        for ip in connected_ips:
            if ip not in allowed_ips:
                messages = "*{}，异常 IP 地址 {} 通过 ssh 访问香港图片服务器，请关注！*".format(
                    current_time, ip
                )
                send_notification(messages, ip)
            else:
                info_message = "{}，{} 无异常。".format(current_time, ip)
                print(info_message)
                logging.info(info_message)
        time.sleep(60)


if __name__ == "__main__":
    main()

